package com.example.vpnservicedemo.receive;

import android.net.VpnService;

import com.example.vpnservicedemo.ReadVpn;

import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.IpV4Packet;
import org.pcap4j.packet.IpV4Rfc1349Tos;
import org.pcap4j.packet.IpV6Packet;
import org.pcap4j.packet.IpV6SimpleFlowLabel;
import org.pcap4j.packet.IpV6SimpleTrafficClass;
import org.pcap4j.packet.Packet;
import org.pcap4j.packet.TcpPacket;
import org.pcap4j.packet.UnknownPacket;
import org.pcap4j.packet.namednumber.IpNumber;
import org.pcap4j.packet.namednumber.IpVersion;
import org.pcap4j.packet.namednumber.TcpPort;

import java.io.IOException;
import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

public class TcpPipe implements Pipe {

    SocketChannel socketChannel;

    private IpPacket srcIpPacket;
    private VpnService vpnService;

    private ReadVpn readVpn;


    AtomicInteger replySeq = new AtomicInteger(1000);
    AtomicInteger replyAck = new AtomicInteger(0);

    AtomicBoolean sendFin = new AtomicBoolean(false);
    AtomicBoolean replyFin = new AtomicBoolean(false);


    private final Map<Integer, TcpPacket> dataMap = new ConcurrentHashMap<>();


    Thread orderWriteThread;

    public TcpPipe(VpnService vpnService, ReadVpn readVpn) {
        this.vpnService = vpnService;
        this.readVpn = readVpn;

        orderWriteThread = new Thread(new OrderWrite());
        orderWriteThread.setName("OrderWriteThread");
        orderWriteThread.start();

    }





    protected LinkedBlockingQueue<byte[]> replyPshQueue = new LinkedBlockingQueue<>();

    class OrderWrite implements Runnable{

        @Override
        public void run() {
            while (!Thread.interrupted()) {
                try {

                    byte[] take = replyPshQueue.take();

                    if(take == null){

                        if (!sendFin.getAndSet(true)) {
                            replyFin();
                        }

                    }else{
                        replyPsh(take);
                    }


                } catch (InterruptedException e) {
                    e.printStackTrace();
                }

            }
        }
    }




    @Override
    public int receive(IpPacket ipPacket) throws IOException {

        TcpPacket tcpPacket = (TcpPacket) ipPacket.getPayload();

        boolean syn = tcpPacket.getHeader().getSyn();
        boolean fin = tcpPacket.getHeader().getFin();
        boolean psh = tcpPacket.getHeader().getPsh();
        boolean ack = tcpPacket.getHeader().getAck();
        boolean rst = tcpPacket.getHeader().getRst();

        InetAddress dstAddr = ipPacket.getHeader().getDstAddr();
        int dstPort = tcpPacket.getHeader().getDstPort().valueAsInt();

        int sequenceNumber = tcpPacket.getHeader().getSequenceNumber();

        if (rst) {
            close();
            return -1;
        }

        if (sendFin.get() && replyFin.get()) {
            close();
            return -1;
        }

        if (syn) {

            if (socketChannel == null) {

                srcIpPacket = ipPacket;
                replyAck.set(sequenceNumber + 1);

                // 使用nio
                socketChannel = SocketChannel.open();
                // 使连接不被自身拦截
                vpnService.protect(socketChannel.socket());
                // 使用非阻塞io
                socketChannel.configureBlocking(false);
                socketChannel.connect(new InetSocketAddress(dstAddr, dstPort));

                //注册事件
                TcpSelector.registerConnect(socketChannel, this);

            }

        } else {

            if (socketChannel != null && !socketChannel.isConnected()) {
                //连接重置
                replyRst(tcpPacket);
                return -1;
            }

            try {
                Packet payload = tcpPacket.getPayload();

                if ((payload != null && payload.getRawData().length > 0) || fin) {

                    //刚好到发送序列只发送
                    if (replyAck.get() == sequenceNumber) {

                        sendData(tcpPacket);

                        while (true) {

                            TcpPacket packets = dataMap.get(replyAck.get());
                            if (packets != null) {
                                dataMap.remove(replyAck.get());
                                sendData(packets);

                            } else {
                                break;
                            }
                        }

                    } else {
                        dataMap.put(sequenceNumber, tcpPacket);
                    }

                }

            } catch (Exception e) {
                e.printStackTrace();
                System.out.println("sendData 异常:" + e.getMessage());
                close();
                return -1;
            }

        }

        return 0;
    }

    private void sendData(TcpPacket tcpPacket) throws IOException {
        Packet payload = tcpPacket.getPayload();
        int dataLen = 0;
        if (payload != null && payload.getRawData().length > 0) {
            dataLen = payload.getRawData().length;
        }

        if (dataLen > 0) {

            if (socketChannel.isConnected()) {
                socketChannel.write(ByteBuffer.wrap(payload.getRawData()));
            }


        }

        int add = dataLen + (tcpPacket.getHeader().getFin() ? 1 : 0);

        replyAck(replyAck.get() + add);

        replyAck.addAndGet(add);

        if (tcpPacket.getHeader().getFin()) {
            replyFin.set(true);
        }

        if (tcpPacket.getHeader().getFin()) {
            if (sendFin.getAndSet(true)) {
                replyFin();
            }
        }


    }

    private void writeData(byte[] bytes) {
        readVpn.pushVpn(bytes);
    }

    protected void close() {
        if (socketChannel != null) {
            try {
                socketChannel.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        if(orderWriteThread!=null){
            orderWriteThread.interrupt();
        }
    }

    private void replyRst(TcpPacket tcpPacket) {
        int sequenceNumber = tcpPacket.getHeader().getSequenceNumber();

        if (tcpPacket.getPayload() != null) {
            sequenceNumber += tcpPacket.getPayload().getRawData().length;
        }

        replyAck.set(sequenceNumber);
        IpPacket ipPack = buildIpPack(true);
        byte[] rawData = ipPack.getRawData();
        writeData(rawData);

    }


    public void replyKeepAlive() {

        IpPacket newIpacket = buildIpPack(false, false, false, true, false, null, replySeq.get() - 1, replyAck.get());
        byte[] rawData = newIpacket.getRawData();

        writeData(rawData);
    }

    public void replySyn() {

        IpPacket newIpacket = buildIpPack(true, false, false, true, null);
        byte[] rawData = newIpacket.getRawData();
        writeData(rawData);

        replySeq.incrementAndGet();
    }

    public void replyFin() {

        int seqNumber = replySeq.get();
        int ackNumber = replyAck.get();
        System.out.println("seq:" + seqNumber + ",ack:" + ackNumber);

        IpPacket newIpacket = buildIpPack(false, false, true, true, false, null, seqNumber, ackNumber);
        byte[] rawData = newIpacket.getRawData();
        writeData(rawData);


        replySeq.set(seqNumber + 1);
    }

    public void replyAck() {

        IpPacket newIpacket = buildIpPack(false, false, false, true, null);
        byte[] rawData = newIpacket.getRawData();

        writeData(rawData);
    }

    public void replyAck(int ack) {

        IpPacket newIpacket = buildIpPack(false, false, false, true, false, null, replySeq.get(), ack);
        byte[] rawData = newIpacket.getRawData();

        writeData(rawData);
    }

    public void replyPsh(int seq, byte[] data) {

        int ackNumber = replyAck.get();

        IpPacket newIpacket = buildIpPack(false, true, false, true, false, data, seq, ackNumber);
        byte[] rawData = newIpacket.getRawData();

        writeData(rawData);

        //this.replySeq.getAndAdd(data.length);

    }

    public void replyPsh(byte[] data) {

        IpPacket newIpacket = buildIpPack(false, true, false, true, data);
        byte[] rawData = newIpacket.getRawData();

        writeData(rawData);

        this.replySeq.getAndAdd(data.length);

    }


    private IpPacket buildIpPack(boolean syn, boolean psh, boolean fin, boolean ack, byte[] data) {
        int seqNumber = replySeq.get();
        int ackNumber = replyAck.get();
        return buildIpPack(syn, psh, fin, ack, false, data, seqNumber, ackNumber);
    }

    private IpPacket buildIpPack(boolean rst) {
        int seqNumber = replySeq.get();
        int ackNumber = replyAck.get();
        return buildIpPack(false, false, false, false, true, null, seqNumber, ackNumber);
    }


    private IpPacket buildIpPack(boolean syn, boolean psh, boolean fin, boolean ack, boolean rst, byte[] data, int seqNumber, int ackNumber) {

        TcpPacket localtcpPacket = (TcpPacket) srcIpPacket.getPayload();
        TcpPort srcPort = localtcpPacket.getHeader().getSrcPort();
        TcpPort dstPort = localtcpPacket.getHeader().getDstPort();


        InetAddress srcAddr = srcIpPacket.getHeader().getSrcAddr();
        InetAddress dstAddr = srcIpPacket.getHeader().getDstAddr();

        TcpPacket.Builder tcpBuilder;
        tcpBuilder = new TcpPacket.Builder()
                .srcPort(dstPort)
                .dstPort(srcPort)
                .sequenceNumber(seqNumber)
                .window(Short.MIN_VALUE)
                .acknowledgmentNumber(ackNumber)
                .syn(syn)
                .psh(psh)
                .fin(fin)
                .ack(ack)
                .rst(rst)
                .srcAddr(dstAddr)
                .dstAddr(srcAddr)
                .correctChecksumAtBuild(true)
                .correctLengthAtBuild(true);


        if (psh) {
            UnknownPacket.Builder unknownBuilder = new UnknownPacket.Builder();
            unknownBuilder.rawData(data);
            tcpBuilder.payloadBuilder(unknownBuilder);
        }


        //创建一个ip包
        IpPacket newIpPacket;
        if (srcIpPacket instanceof IpV6Packet) {
            newIpPacket = new IpV6Packet.Builder()
                    .version(IpVersion.IPV6)
                    .srcAddr((Inet6Address) dstAddr)
                    .dstAddr((Inet6Address) srcAddr)
                    .correctLengthAtBuild(true)
                    .payloadBuilder(tcpBuilder)
                    .trafficClass(IpV6SimpleTrafficClass.newInstance((byte) 0))
                    .flowLabel(IpV6SimpleFlowLabel.newInstance((byte) 0))
                    .nextHeader(IpNumber.IPV6)
                    .build();
        } else {

            newIpPacket = new IpV4Packet.Builder()
                    .version(IpVersion.IPV4)
                    .ttl((byte) 45)
                    .tos(IpV4Rfc1349Tos.newInstance((byte) 0))
                    .protocol(IpNumber.TCP)
                    .srcAddr((Inet4Address) dstAddr)
                    .dstAddr((Inet4Address) srcAddr)
                    .correctLengthAtBuild(true)
                    .correctChecksumAtBuild(true)
                    .payloadBuilder(tcpBuilder)
                    .identification((byte) getIpId())
                    .build();
        }

        return newIpPacket;
    }


    private static AtomicInteger ipId = new AtomicInteger(0);

    public static int getIpId() {
        return ipId.incrementAndGet();
    }


}
